import numpy as np
import roboticstoolbox as rtb
from spatialmath import SE3


class Robot:
    def __init__(self):
        self.a1 = 1.0
        self.a2 = 1.0
        link1 = rtb.DHLink(d=0, alpha=0, a=self.a1, theta=0, m=1, r=np.array((0, 0, 0)))
        link2 = rtb.DHLink(d=0, alpha=0, a=self.a2, theta=0, m=1, r=np.array((0, 0, 0)))

        link = [link1, link2]
        self.robot = rtb.DHRobot(link, base=SE3.Trans(-0.5, 0, 0) * SE3.Rx(np.pi / 2))

    def get_static_torque(self, q, joint: bool = True):
        if joint:
            return self.robot.gravload(q)

        return self.robot.gravload(self.ikine(q))

    def get_dynamic_torque(self, q, dq, ddq, joint: bool = True):
        return self.get_torque(q, dq, ddq, joint) - self.get_static_torque(q, joint)

    def get_torque(self, q, dq, ddq, joint: bool = True):
        if joint:
            return self.robot.rne(q, dq, ddq)
        return self.robot.rne(self.ikine(q), self.get_dq(q, dq), self.get_ddq(q, dq, ddq))

    def plot(self, q):
        self.robot.plot(q, block=True)

    def get_jacob(self, q):
        return self.robot.jacob0(q)

    def get_hessian(self, q):
        return self.robot.hessian0(q)

    def get_dq(self, x, dx):
        q = self.ikine(x)
        J = self.robot.jacob0(q)
        return np.linalg.pinv(J[[0, 2], :]) @ dx

    def get_ddq(self, x, dx, ddx):
        q = self.ikine(x)
        J = self.robot.jacob0(q)
        H = self.robot.hessian0(q)
        return (np.linalg.pinv(J[[0, 2], :])) @ (ddx - ((H @ dx).T @ dx)[[0, 2]])

    def ikine(self, x) -> tuple:
        x = [x[0] - self.robot.base.t[0], x[1] - self.robot.base.t[2]]
        a0 = np.linalg.norm(x)
        q0 = np.arctan2(x[1] / a0, x[0] / a0)

        q1 = np.arccos((x[0] ** 2 + x[1] ** 2 + self.a1 ** 2 - self.a2 ** 2) / (2 * a0 * self.a1)) + q0
        q2 = -np.arccos((x[0] ** 2 + x[1] ** 2 + self.a2 ** 2 - self.a1 ** 2) / (2 * a0 * self.a2)) + q0 - q1
        return q1, q2


if __name__ == '__main__':
    robot = Robot()
    q = [1, 2]
    dq = [0, 0]
    ddq = [0, 0]
    static_torque = robot.get_static_torque(q)
    dynamic_torque = robot.get_dynamic_torque(q, dq, ddq)
    torque = robot.get_torque(q, dq, ddq)
    print('static torque: ', static_torque)
    print('dynamic torque: ', dynamic_torque)
    print('torque: ', torque)

    jacob = robot.get_jacob(q)
    print('jacob: ', jacob)
    hessian = robot.get_hessian(q)
    print('hessian: ', hessian)

    x = (1, 1)
    q = robot.ikine((1, 1))
    print('q: ', q)

    dx = (0, 0)
    dq = robot.get_dq(x, dx)
    print('dq: ', dq)

    ddx = (1.0, 0)
    ddq = robot.get_ddq(x, dx, ddx)
    print('ddq: ', ddq)

